package main

import (
    "adkr/pkg/adkr"
    "adkr/pkg/crypto"
    "encoding/json"
    "flag"
    "fmt"
    "log"
    "os"
    "sync"
    "time"
)

// TestConfig represents a test configuration
type TestConfig struct {
    Name        string `json:"name"`
    FastCoin    bool   `json:"fast_coin"`
    StrongPred  bool   `json:"strong_predicate"`
    ABAEvidence bool   `json:"aba_evidence"`
    KRSPBC      bool   `json:"kr_spbc"`
    RBCMode     string `json:"rbc_mode"`
    CoinSkipVerify bool `json:"coin_skip_verify"`
}

// TestResult represents the result of a test run
type TestResult struct {
    Config      TestConfig `json:"config"`
    Success     bool       `json:"success"`
    Duration    int64      `json:"duration_ms"`
    Decision    string     `json:"decision,omitempty"`
    Error       string     `json:"error,omitempty"`
    MessagesSent []int64   `json:"messages_sent"`
    MessagesRecv []int64   `json:"messages_received"`
    Timeouts    int        `json:"timeouts"`
    CoinMetrics map[string]interface{} `json:"coin_metrics,omitempty"`
}

// ParityTest runs comprehensive parity tests
type ParityTest struct {
    n, f, k int
    results []TestResult
}

func NewParityTest(n, f, k int) *ParityTest {
    return &ParityTest{
        n: n, f: f, k: k,
        results: make([]TestResult, 0),
    }
}

// RunTest executes a single test configuration
func (pt *ParityTest) RunTest(config TestConfig) TestResult {
    log.Printf("Running test: %s", config.Name)
    
    start := time.Now()
    result := TestResult{
        Config: config,
        Success: false,
    }
    
    // Set global flags
    adkr.SetFastBench(config.FastCoin)
    adkr.SetStrongMVBAEnabled(config.StrongPred)
    adkr.SetABAEvidence(config.ABAEvidence)
    adkr.SetKeyRefreshSPBCFast(config.KRSPBC)
    adkr.SetRBCMode(config.RBCMode)
    adkr.SetCoinSkipVerify(config.CoinSkipVerify)
    
    // Create test environment
    sid := fmt.Sprintf("parity-test-%d", time.Now().UnixNano())
    
    // Generate cryptographic materials
    g1 := crypto.GetG1()
    g2 := crypto.GetG2()
    epks := make([]crypto.G1, pt.n)
    esks := make([]crypto.ZR, pt.n)
    for i := 0; i < pt.n; i++ {
        esks[i] = crypto.ZRFromInt(int64(1000 + i))
        epks[i] = crypto.MulG1(&g1, &esks[i])
    }
    
    // Generate threshold BLS keys
    thpk, _, thsks := crypto.GenerateThresholdKeys(pt.n, pt.f+1)
    thpks := make([]crypto.G1, pt.n)
    for i := 0; i < pt.n; i++ {
        thpks[i] = crypto.MulG1(&g1, &thsks[i])
    }
    
    // Create message bus for simulation
    bus := newTestBus(pt.n)
    defer bus.stop()
    
    // Launch protocol instances
    var wg sync.WaitGroup
    outputs := make([]chan adkr.ADKRHighOutput, pt.n)
    for i := 0; i < pt.n; i++ {
        outputs[i] = make(chan adkr.ADKRHighOutput, 1)
    }
    
    for pid := 0; pid < pt.n; pid++ {
        wg.Add(1)
        pid := pid
        go func() {
            defer wg.Done()
            
            cfg := adkr.ADKRHighConfig{
                SID:   sid,
                PID:   pid,
                N:     pt.n,
                F:     pt.f,
                K:     pt.k,
                G1:    &g1,
                G2:    &g2,
                EPKs:  epks,
                ESK:   esks[pid],
                THPK:  *thpk,
                THPKs: thpks,
                THSK:  thsks[pid],
            }
            
            proto := adkr.NewADKRHighProtocol(cfg, bus.ins[pid], bus.outs[pid], outputs[pid], log.Default())
            
            if err := proto.Run(); err != nil {
                result.Error = fmt.Sprintf("PID %d: %v", pid, err)
                return
            }
            
            select {
            case out := <-outputs[pid]:
                // Extract decision from config chain
                if len(out.ConfigChain) > 0 {
                    result.Decision = fmt.Sprintf("%v", out.ConfigChain[len(out.ConfigChain)-1])
                }
                result.Success = true
            case <-time.After(30 * time.Second):
                result.Error = fmt.Sprintf("PID %d: timeout", pid)
            }
        }()
    }
    
    wg.Wait()
    result.Duration = time.Since(start).Milliseconds()
    
    // Collect message statistics
    sentCnt, recvCnt := bus.snapshot()
    result.MessagesSent = sentCnt
    result.MessagesRecv = recvCnt
    
    log.Printf("Test %s completed: success=%v, duration=%dms", config.Name, result.Success, result.Duration)
    return result
}

// RunAllTests executes all test configurations
func (pt *ParityTest) RunAllTests() {
    testConfigs := []TestConfig{
        // Fast coin tests
        {"fast-coin-weak-pred-classic", true, false, false, false, "", false},
        {"fast-coin-weak-pred-dual", true, false, false, false, "dual", false},
        {"fast-coin-weak-pred-dual-cp", true, false, false, false, "dual_cp", false},
        {"fast-coin-strong-pred-classic", true, true, false, false, "", false},
        {"fast-coin-strong-pred-dual", true, true, false, false, "dual", false},
        {"fast-coin-strong-pred-dual-cp", true, true, false, false, "dual_cp", false},
        
        // Real coin tests
        {"real-coin-weak-pred-classic", false, false, false, false, "", false},
        {"real-coin-weak-pred-dual", false, false, false, false, "dual", false},
        {"real-coin-weak-pred-dual-cp", false, false, false, false, "dual_cp", false},
        {"real-coin-strong-pred-classic", false, true, false, false, "", false},
        {"real-coin-strong-pred-dual", false, true, false, false, "dual", false},
        {"real-coin-strong-pred-dual-cp", false, true, true, false, "dual_cp", false},
        
        // SPBC fast-path tests
        {"fast-coin-spbc-weak-pred", true, false, false, true, "", false},
        {"fast-coin-spbc-strong-pred", true, true, false, true, "", false},
        {"real-coin-spbc-weak-pred", false, false, false, true, "", false},
        {"real-coin-spbc-strong-pred", false, true, false, true, "", false},
        
        // ABA evidence tests
        {"fast-coin-aba-evidence", true, false, true, false, "", false},
        {"real-coin-aba-evidence", false, false, true, false, "", false},
        
        // Coin skip verify tests
        {"fast-coin-skip-verify", true, false, false, false, "", true},
        {"real-coin-skip-verify", false, false, false, false, "", true},
    }
    
    log.Printf("Running %d parity tests...", len(testConfigs))
    
    for _, config := range testConfigs {
        result := pt.RunTest(config)
        pt.results = append(pt.results, result)
        
        // Small delay between tests
        time.Sleep(100 * time.Millisecond)
    }
}

// GenerateReport creates a comprehensive test report
func (pt *ParityTest) GenerateReport() {
    log.Printf("Generating parity test report...")
    
    // Calculate statistics
    totalTests := len(pt.results)
    successfulTests := 0
    totalDuration := int64(0)
    
    decisions := make(map[string]int)
    
    for _, result := range pt.results {
        if result.Success {
            successfulTests++
            totalDuration += result.Duration
            decisions[result.Decision]++
        }
    }
    
    // Create summary
    summary := map[string]interface{}{
        "total_tests": totalTests,
        "successful_tests": successfulTests,
        "success_rate": float64(successfulTests) / float64(totalTests),
        "average_duration_ms": float64(totalDuration) / float64(successfulTests),
        "decision_distribution": decisions,
        "test_results": pt.results,
    }
    
    // Output JSON report
    jb, _ := json.MarshalIndent(summary, "", "  ")
    fmt.Println(string(jb))
    
    log.Printf("Parity test completed: %d/%d tests successful (%.1f%%)", 
        successfulTests, totalTests, float64(successfulTests)/float64(totalTests)*100)
}

// testBus simulates network communication for testing
type testBus struct {
    n       int
    ins     []chan adkr.Message
    outs    []chan adkr.Message
    sentCnt []int64
    recvCnt []int64
    stopCh  chan struct{}
    wg      sync.WaitGroup
}

func newTestBus(n int) *testBus {
    bus := &testBus{
        n:       n,
        ins:     make([]chan adkr.Message, n),
        outs:    make([]chan adkr.Message, n),
        sentCnt: make([]int64, n),
        recvCnt: make([]int64, n),
        stopCh:  make(chan struct{}),
    }
    
    for i := 0; i < n; i++ {
        bus.ins[i] = make(chan adkr.Message, 100)
        bus.outs[i] = make(chan adkr.Message, 100)
    }
    
    // Start message routing
    bus.wg.Add(1)
    go func() {
        defer bus.wg.Done()
        for {
            select {
            case <-bus.stopCh:
                return
            default:
                // Route messages between nodes
                for i := 0; i < n; i++ {
                    select {
                    case msg := <-bus.outs[i]:
                        // Broadcast to all other nodes
                        for j := 0; j < n; j++ {
                            if j != i {
                                select {
                                case bus.ins[j] <- msg:
                                    bus.recvCnt[j]++
                                default:
                                }
                            }
                        }
                        bus.sentCnt[i]++
                    default:
                    }
                }
                time.Sleep(1 * time.Millisecond)
            }
        }
    }()
    
    return bus
}

func (b *testBus) stop() {
    close(b.stopCh)
    b.wg.Wait()
    for i := 0; i < b.n; i++ {
        close(b.ins[i])
        close(b.outs[i])
    }
}

func (b *testBus) snapshot() (sent []int64, recv []int64) {
    sent = make([]int64, b.n)
    recv = make([]int64, b.n)
    for i := 0; i < b.n; i++ {
        sent[i] = b.sentCnt[i]
        recv[i] = b.recvCnt[i]
    }
    return
}

func main() {
    n := flag.Int("n", 4, "Number of nodes")
    f := flag.Int("f", 1, "Fault tolerance")
    k := flag.Int("k", 1, "Number of rounds")
    output := flag.String("output", "", "Output file for results (default: stdout)")
    flag.Parse()
    
    log.Printf("Starting parity tests: n=%d, f=%d, k=%d", *n, *f, *k)
    
    pt := NewParityTest(*n, *f, *k)
    pt.RunAllTests()
    
    if *output != "" {
        // Write results to file
        summary := map[string]interface{}{
            "test_config": map[string]int{"n": *n, "f": *f, "k": *k},
            "results": pt.results,
        }
        jb, _ := json.MarshalIndent(summary, "", "  ")
        if err := os.WriteFile(*output, jb, 0644); err != nil {
            log.Fatalf("Failed to write results: %v", err)
        }
        log.Printf("Results written to %s", *output)
    } else {
        pt.GenerateReport()
    }
}



